#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
TCP 转发模块 - 处理 TCP 连接转发
"""

import socket
import selectors
import logging
import errno
import threading
import time
from connection_pool import ConnectionPool
from logger import setup_logger
from security import is_ip_allowed
from config import get_config_value

# 设置日志记录器
logger = setup_logger("tcp_forwarder")

def tcp_forward(local_port, remote_host, remote_port):
    """TCP 端口转发 - 使用 selectors 优化和连接池"""
    try:
        # 创建连接池，增加最大连接数
        connection_pool = ConnectionPool(remote_host, remote_port)
        
        # 启动定时统计日志
        def log_stats_periodically():
            while True:
                connection_pool.log_stats()
                time.sleep(60)  # 每分钟记录一次
        
        stats_thread = threading.Thread(target=log_stats_periodically, daemon=True)
        stats_thread.start()
        
        # 创建服务器套接字
        server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        # 设置接收和发送缓冲区大小
        server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF, 262144)
        server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF, 262144)
        
        # 启用TCP自动窗口缩放
        if hasattr(socket, 'TCP_AUTOTUNING'):  # Windows Vista及以上支持
            server_socket.setsockopt(socket.IPPROTO_TCP, socket.TCP_AUTOTUNING, 1)
        
        # 设置非阻塞模式
        server_socket.setblocking(False)
        server_socket.bind(('0.0.0.0', local_port))
        server_socket.listen(100)  # 增加监听队列大小
        logger.info(f"TCP 转发已启动: 0.0.0.0:{local_port} -> {remote_host}:{remote_port}")
        
        # 创建 selector
        sel = selectors.DefaultSelector()
        sel.register(server_socket, selectors.EVENT_READ, 
                    lambda s, sel, conns, *args: accept_tcp_connection(s, sel, conns, connection_pool))
        
        # 存储连接信息
        connections = {}
        
        # 事件循环
        while True:
            events = sel.select(timeout=1)
            for key, mask in events:
                callback = key.data
                callback(key.fileobj, sel, connections, None, None)
                
    except Exception as e:
        logger.error(f"TCP 转发错误: {e}", exc_info=True)
    finally:
        if 'connection_pool' in locals():
            connection_pool.cleanup()
        if 'server_socket' in locals():
            server_socket.close()
        if 'sel' in locals():
            sel.close()

def accept_tcp_connection(sock, sel, connections, connection_pool):
    """接受新的 TCP 连接"""
    client_socket, addr = sock.accept()
    
    # 检查客户端 IP 是否允许连接
    if not is_ip_allowed(addr[0]):
        logger.warning(f"拒绝来自 {addr[0]}:{addr[1]} 的 TCP 连接 - IP 不在允许列表中")
        client_socket.close()
        return
    
    client_socket.setblocking(False)
    # 设置 SO_KEEPALIVE 选项
    client_socket.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)
    # 设置接收和发送缓冲区大小
    client_socket.setsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF, 262144)  # 增加到256KB
    client_socket.setsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF, 262144)  # 增加到256KB
    # 禁用 Nagle 算法，减少小数据包的延迟
    client_socket.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
    
    # 移除不支持的TCP窗口大小设置
    # client_socket.setsockopt(socket.SOL_TCP, socket.TCP_WINDOW_CLAMP, 262144)
    
    # 在 Windows 上设置 TCP keepalive 参数
    client_socket.ioctl(socket.SIO_KEEPALIVE_VALS, (1, 60000, 5000))
    
    logger.info(f"接受来自 {addr[0]}:{addr[1]} 的 TCP 连接")
    
    # 从连接池获取远程连接
    try:
        remote_socket = connection_pool.get_connection()
        remote_socket.setblocking(False)
        
        # 注册客户端和远程套接字
        sel.register(client_socket, selectors.EVENT_READ, 
                    lambda s, sel, conns, *args: handle_tcp_data(s, sel, conns, True, connection_pool))
        sel.register(remote_socket, selectors.EVENT_READ, 
                    lambda s, sel, conns, *args: handle_tcp_data(s, sel, conns, False, connection_pool))
        
        # 存储连接信息
        connections[client_socket] = remote_socket
        connections[remote_socket] = client_socket
    except Exception as e:
        logger.error(f"获取远程连接失败: {e}", exc_info=True)
        client_socket.close()

def handle_tcp_data(sock, sel, connections, is_client, connection_pool):
    """处理 TCP 数据传输"""
    try:
        data = sock.recv(get_config_value("tcp", "buffer_size", 8192))  # 增加缓冲区大小
        if not data:
            # 连接关闭
            peer_sock = connections[sock]
            del connections[peer_sock]
            del connections[sock]
            sel.unregister(sock)
            sel.unregister(peer_sock)
            
            # 如果是客户端连接关闭，将远程连接释放回连接池
            if is_client:
                connection_pool.release_connection(peer_sock)
                logger.debug("客户端连接关闭，远程连接已释放回连接池")
            else:
                # 如果是远程连接关闭，关闭客户端连接
                sock.close()
                peer_sock.close()
                logger.debug("远程连接关闭，客户端连接已关闭")
            return
            
        # 转发数据
        peer_sock = connections[sock]
        try:
            peer_sock.send(data)
        except BlockingIOError:
            # 发送缓冲区已满，这是正常的非阻塞行为
            if logger.isEnabledFor(logging.DEBUG):
                logger.debug("发送缓冲区已满，稍后再试")
            return
        
        # 记录数据传输（调试级别）
        if logger.isEnabledFor(logging.DEBUG):
            direction = "客户端->远程" if is_client else "远程->客户端"
            logger.debug(f"数据传输: {direction}, 大小: {len(data)} 字节")
            
    except BlockingIOError:
        # 接收缓冲区为空，这是正常的非阻塞行为
        return
    except socket.error as e:
        # 处理常见的套接字错误
        if e.errno == errno.WSAEWOULDBLOCK:
            # 等同于 BlockingIOError，忽略
            return
        elif e.errno == errno.WSAECONNRESET:
            # 连接被重置，记录但不需要详细错误信息
            logger.info(f"TCP 连接被重置: {e}")
        else:
            # 其他套接字错误，记录详细信息
            logger.error(f"TCP 数据处理错误: {e}")
        
        try:
            # 清理连接资源
            if sock in connections:
                peer_sock = connections[sock]
                
                # 从连接字典中删除
                if peer_sock in connections:
                    del connections[peer_sock]
                del connections[sock]
                
                # 从选择器中注销
                try:
                    sel.unregister(sock)
                except Exception:
                    pass
                
                try:
                    sel.unregister(peer_sock)
                except Exception:
                    pass
                
                # 如果是客户端连接出错，将远程连接释放回连接池
                if is_client:
                    connection_pool.release_connection(peer_sock)
                    logger.debug("客户端连接出错，远程连接已释放回连接池")
                else:
                    # 如果是远程连接出错，关闭客户端连接
                    sock.close()
                    peer_sock.close()
                    logger.debug("远程连接出错，客户端连接已关闭")
            else:
                # 如果连接不在字典中，直接关闭
                sock.close()
        except Exception as inner_e:
            logger.error(f"清理连接时出错: {inner_e}")